package com.cashop.mcp.sse;

import com.cashop.mcp.sse.model.Rule;
import io.vertx.core.Vertx;
import io.vertx.core.http.HttpServer;
import io.vertx.core.http.HttpServerResponse;
import io.vertx.core.json.JsonArray;
import io.vertx.core.json.JsonObject;
import io.vertx.ext.web.Router;
import io.vertx.ext.web.RoutingContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Stream;

public class Main {
  private static final Logger logger = LoggerFactory.getLogger(Main.class);
  private static Vertx vertx;
  // 存储活跃的SSE连接，key为sessionId，value为对应的response
  private static Map<String, HttpServerResponse> activeConnections = new ConcurrentHashMap<>();
  
  // 规则缓存相关字段
  private static List<Rule> cachedRules = new ArrayList<>();
  private static long lastCacheUpdate = 0;
  private static final long CACHE_TTL = 5000; // 5秒缓存时间

  public static void main(String[] args) {
    vertx = Vertx.vertx();
    Router router = Router.router(vertx);
    
    // 创建SSE端点
    router.get("/sse").handler(Main::handleSSE);
    
    // 创建处理JSON-RPC请求的POST端点
    router.post("/messages").handler(Main::handleMessages);
    
    // 添加兜底处理程序，打印所有入参和请求头
    router.route("/*").handler(Main::handleFallback);
    
    // 添加定时任务，定期拉取rules目录的最新代码
    setupGitPullTask();
    
    HttpServer server = vertx.createHttpServer();
    server.requestHandler(router)
      .listen(8080, result -> {
        if (result.succeeded()) {
          logger.info("Server started on port 8080");
        } else {
          logger.error("Failed to start server: {}", result.cause().getMessage());
        }
      });
  }
  
  private static void handleSSE(RoutingContext context) {
    HttpServerResponse response = context.response();
    response.putHeader("Content-Type", "text/event-stream")
      .putHeader("Connection", "keep-alive")
      .putHeader("Cache-Control", "no-cache")
      .setChunked(true);
    
    // 生成唯一的sessionId
    String sessionId = UUID.randomUUID().toString();
    
    // 发送endpoint事件
    String endpointData = "/messages?sessionId=" + sessionId;
    response.write("event: endpoint\ndata: " + endpointData + "\n\n");
    
    // 将连接存储到活跃连接映射中
    activeConnections.put(sessionId, response);
    
    // 当客户端断开连接时清理资源
    response.closeHandler(v -> {
      activeConnections.remove(sessionId);
      logger.info("Client with sessionId {} disconnected", sessionId);
    });
  }
  
  private static void handleMessages(RoutingContext context) {
    // 立即返回202状态码
    context.response()
      .setStatusCode(202)
      .end();
    
    // 在后台处理请求
    context.request().bodyHandler(buffer -> {
      try {
        // 解析JSON-RPC请求
        JsonObject jsonRpcRequest = new JsonObject(buffer.toString());
        String method = jsonRpcRequest.getString("method");
        String sessionId = context.request().getParam("sessionId");
        String protocolVersion = jsonRpcRequest.getString("protocolVersion");
        int id = jsonRpcRequest.getInteger("id", 1);
        
        logger.info("Received JSON-RPC request: {}", jsonRpcRequest.encodePrettily());

        if ("initialize".equals(method) && sessionId != null) {
          handleInitialize(context, sessionId, protocolVersion, id);
        } else if ("tools/list".equals(method)) {
          handleToolsList(context, id);
        } else if ("tools/call".equals(method)) {
          handleToolsCall(context, jsonRpcRequest, id);
        } else if ("resources/list".equals(method)) {
          handleResourcesList(context, id);
        } else if ("resources/templates/list".equals(method)) {
          handleResourcesTemplatesList(context, id);
        } else if ("ping".equals(method)) {
          handlePing(context, id);
        } else {
          logger.warn("Unsupported method: {}", method);
          printRequestDetails(context);
        }
      } catch (Exception e) {
        logger.error("Error processing request", e);
      }
    });
  }
  
  // 兜底处理程序，打印所有入参和请求头
  private static void handleFallback(RoutingContext context) {
    logger.info("=== Fallback handler triggered ===");
    printRequestDetails(context);
    
    // 返回404状态码
    context.response()
      .setStatusCode(404)
      .end("Not Found");
  }
  
  // 打印请求详情的辅助方法
  private static void printRequestDetails(RoutingContext context) {
    logger.info("Request Path: {}", context.request().path());
    logger.info("Request Method: {}", context.request().method());
    logger.info("Request Headers:");
    context.request().headers().forEach(header -> 
      logger.info("  {}: {}", header.getKey(), header.getValue()));
    
    logger.info("Query Parameters:");
    context.request().params().forEach(param -> 
      logger.info("  {}: {}", param.getKey(), param.getValue()));
  }
  
  private static void handleInitialize(RoutingContext context, String sessionId, String protocolVersion, int id) {
    try {
      // 查找对应的SSE连接
      HttpServerResponse sseResponse = activeConnections.get(sessionId);
      
      if (sseResponse != null) {
        // 构造JSON-RPC响应
        JsonObject responseBody = new JsonObject()
          .put("jsonrpc", "2.0")
          .put("id", id)
          .put("result", new JsonObject()
            .put("protocolVersion", protocolVersion != null ? protocolVersion : "2024-11-05")
            .put("capabilities", new JsonObject()
              .put("tools", new JsonObject().put("listChanged", true))
              .put("resources", new JsonObject().put("subscribe", true).put("listChanged", true))
              .put("prompts", new JsonObject().put("listChanged", true)))
            .put("serverInfo", new JsonObject()
              .put("name", "tollge-mcp")
              .put("version", "1.2.3")));
        
        // 通过SSE连接发送message事件
        sseResponse.write("event: message\ndata: " + responseBody.encode() + "\n\n");
      } else {
        logger.warn("Session not found: {}", sessionId);
        // 当session not found时，告诉client重连
        JsonObject errorBody = new JsonObject()
          .put("jsonrpc", "2.0")
          .put("id", id)
          .put("error", new JsonObject()
            .put("code", -32001)
            .put("message", "Session not found, please reconnect"));
        context.response()
          .setStatusCode(400)
          .end(errorBody.encode());
      }
    } catch (Exception e) {
      logger.error("Error handling initialize request", e);
    }
  }
  
  private static void handleToolsList(RoutingContext context, int id) {
    try {
      JsonArray toolsArray = new JsonArray();

      JsonObject inputSchema = new JsonObject();
      inputSchema.put("type", "object")
              .put("properties", new JsonObject()
                      .put("ruleName", new JsonObject().put("type", "string").put("description", "规则名称"))
              )
              .put("required", new JsonArray().add("ruleName"));
      JsonObject toolObj = new JsonObject()
        .put("name", "read_rule")
        .put("description", "读取规则")
        .put("inputSchema", inputSchema);
//      toolsArray.add(toolObj);

      // list_rules tool
      inputSchema = new JsonObject();
      inputSchema.put("type", "object")
              .put("properties", new JsonObject()
                      .put("keyword", new JsonObject().put("type", "string").put("description", "搜索关键字，用于匹配规则名称和描述")));
      toolObj = new JsonObject()
              .put("name", "list_rules")
              .put("description", "列出规则")
              .put("inputSchema", inputSchema);
      toolsArray.add(toolObj);

      // read_rules tool
      inputSchema = new JsonObject();
      inputSchema.put("type", "object")
              .put("properties", new JsonObject()
                      .put("ruleNameList", new JsonObject().put("type", "array").put("items", "string").put("description", "规则名称列表"))
              )
              .put("required", new JsonArray().add("ruleNameList"));
      toolObj = new JsonObject()
              .put("name", "read_rules")
              .put("description", "读取规则列表")
              .put("inputSchema", inputSchema);
      toolsArray.add(toolObj);

      JsonObject responseBody = new JsonObject()
        .put("jsonrpc", "2.0")
        .put("id", id)
        .put("result", new JsonObject()
          .put("tools", toolsArray));
      
      // 发送消息给客户端
      sendMessageToClient(context, responseBody, id);
    } catch (Exception e) {
      System.err.println("Error handling tools/list request: " + e.getMessage());
      e.printStackTrace();
    }
  }
  
  private static void handleToolsCall(RoutingContext context, JsonObject request, int id) {
    try {
      JsonObject params = request.getJsonObject("params");
      String toolName = params.getString("name");
      JsonObject arguments = params.getJsonObject("arguments");
      
      JsonObject responseBody = new JsonObject();
      
      switch (toolName) {
        case "read_rule":
          handleReadRule(arguments, responseBody, id);
          break;
        case "list_rules":
          handleListRules(arguments, responseBody, id);
          break;
        case "read_rules":
          handleReadRules(arguments, responseBody, id);
          break;
        default:
          responseBody.put("jsonrpc", "2.0")
            .put("id", id)
            .put("error", new JsonObject()
              .put("code", -32601)
              .put("message", "Tool not found: " + toolName));
      }
      
      // 发送消息给客户端
      sendMessageToClient(context, responseBody, id);
    } catch (Exception e) {
      logger.error("Error handling tools/call request", e);
    }
  }
  
  private static void handleReadRule(JsonObject arguments, JsonObject responseBody, int id) {
    try {
      String ruleName = arguments.getString("ruleName");
      List<Rule> rules = scanRules();
      
      for (Rule rule : rules) {
        if (rule.getName().equals(ruleName)) {
          JsonArray contentArray = new JsonArray();
          JsonObject content = new JsonObject()
            .put("type", "text")
            .put("text", rule.getText());
          contentArray.add(content);
          
          responseBody.put("jsonrpc", "2.0")
            .put("id", id)
            .put("result", new JsonObject().put("content", contentArray));
          return;
        }
      }
      
      // 如果没找到规则
      responseBody.put("jsonrpc", "2.0")
        .put("id", id)
        .put("error", new JsonObject()
          .put("code", -32601)
          .put("message", "Rule not found: " + ruleName));
          
    } catch (Exception e) {
      logger.error("Error handling read_rule request", e);
      
      responseBody.put("jsonrpc", "2.0")
        .put("id", id)
        .put("error", new JsonObject()
          .put("code", -32000)
          .put("message", "Internal error: " + e.getMessage()));
    }
  }
  
  private static void handleListRules(JsonObject arguments, JsonObject responseBody, int id) {
    try {
      List<Rule> rules = scanRules();
      JsonArray rulesArray = new JsonArray();
      
      // 获取搜索关键字，如果没有提供则为null
      String keyword = arguments != null ? arguments.getString("keyword") : null;
      
      for (Rule rule : rules) {
        // 如果没有提供关键字，或者关键字匹配规则名称或描述，则添加到结果中
        if (keyword == null ||
            rule.getName().toLowerCase().contains(keyword.toLowerCase()) ||
            rule.getDescription().toLowerCase().contains(keyword.toLowerCase())) {
          JsonObject ruleObj = new JsonObject()
            .put("name", rule.getName())
            .put("description", rule.getDescription());
          rulesArray.add(ruleObj);
        }
      }

      JsonObject content = new JsonObject()
        .put("type", "text")
        .put("text", rulesArray.toString());

      responseBody.put("jsonrpc", "2.0")
        .put("id", id)
        .put("result", new JsonObject().put("content", new JsonArray().add(content)));
        
    } catch (Exception e) {
      logger.error("Error handling list_rules request", e);
      
      responseBody.put("jsonrpc", "2.0")
        .put("id", id)
        .put("error", new JsonObject()
          .put("code", -32000)
          .put("message", "Internal error: " + e.getMessage()));
    }
  }
  
  private static void handleReadRules(JsonObject arguments, JsonObject responseBody, int id) {
    try {
      JsonArray ruleNameList = arguments.getJsonArray("ruleNameList");
      List<Rule> rules = scanRules();
      JsonArray rulesArray = new JsonArray();
      
      for (int i = 0; i < ruleNameList.size(); i++) {
        String ruleName = ruleNameList.getString(i);
        
        for (Rule rule : rules) {
          if (rule.getName().equals(ruleName)) {
            JsonObject ruleObj = new JsonObject()
              .put("name", rule.getName())
              .put("description", rule.getDescription())
              .put("text", rule.getText());
            rulesArray.add(ruleObj);
            break;
          }
        }
      }

      JsonObject content = new JsonObject()
              .put("type", "text")
              .put("text", rulesArray.toString());
      
      responseBody.put("jsonrpc", "2.0")
        .put("id", id)
        .put("result", new JsonObject().put("content", new JsonArray().add(content)));
        
    } catch (Exception e) {
      logger.error("Error handling read_rules request", e);
      
      responseBody.put("jsonrpc", "2.0")
        .put("id", id)
        .put("error", new JsonObject()
          .put("code", -32000)
          .put("message", "Internal error: " + e.getMessage()));
    }
  }

  private static void handleResourcesList(RoutingContext context, int id) {
    try {
      // 创建一个空的资源列表
      JsonArray resourcesArray = new JsonArray();
      
      JsonObject responseBody = new JsonObject()
        .put("jsonrpc", "2.0")
        .put("id", id)
        .put("result", new JsonObject()
          .put("resources", resourcesArray));
      
      // 发送消息给客户端
      sendMessageToClient(context, responseBody, id);
    } catch (Exception e) {
      logger.error("Error handling resources/list request", e);
    }
  }
  
  private static void handleResourcesTemplatesList(RoutingContext context, int id) {
    try {
      // 创建一个空的资源模板列表
      JsonArray resourcesTemplatesArray = new JsonArray();
      
      JsonObject responseBody = new JsonObject()
        .put("jsonrpc", "2.0")
        .put("id", id)
        .put("result", new JsonObject()
          .put("resources", resourcesTemplatesArray));
      
      // 发送消息给客户端
      sendMessageToClient(context, responseBody, id);
    } catch (Exception e) {
      logger.error("Error handling resources/templates/list request", e);
    }
  }
  
  private static void handlePing(RoutingContext context, int id) {
    try {
      // 构造JSON-RPC响应
      JsonObject responseBody = new JsonObject()
        .put("jsonrpc", "2.0")
        .put("id", id)
        .put("result", "pong");
      
      // 发送消息给客户端
      sendMessageToClient(context, responseBody, id);
    } catch (Exception e) {
      logger.error("Error handling ping request", e);
    }
  }
  
  /**
   * 发送消息给客户端的通用方法
   * @param context 路由上下文
   * @param responseBody 响应体
   * @param id 请求ID
   */
  private static void sendMessageToClient(RoutingContext context, JsonObject responseBody, int id) {
    // 查找对应的SSE连接（通过查询参数中的sessionId）
    String sessionId = context.request().getParam("sessionId");
    if (sessionId != null) {
      HttpServerResponse sseResponse = activeConnections.get(sessionId);
      if (sseResponse != null) {
        sseResponse.write("event: message\ndata: " + responseBody.encode() + "\n\n");
      } else {
        // 打印日志
        String methodName = new Throwable().getStackTrace()[2].getMethodName();
        logger.warn("Session not found for {}: {}", methodName, sessionId);
        // 当session not found时，告诉client重连
        JsonObject errorBody = new JsonObject()
          .put("jsonrpc", "2.0")
          .put("id", id)
          .put("error", new JsonObject()
            .put("code", -32001)
            .put("message", "Session not found, please reconnect"));
        context.response()
          .setStatusCode(400)
          .end(errorBody.encode());
      }
    }
  }
  
  private static List<Rule> scanRules() throws IOException {
    long currentTime = System.currentTimeMillis();
    
    // 检查是否需要更新缓存
    if (cachedRules.isEmpty() || (currentTime - lastCacheUpdate) > CACHE_TTL) {
      // 更新缓存
      cachedRules = loadRulesFromDisk();
      lastCacheUpdate = currentTime;
    }
    
    return cachedRules;
  }
  
  private static List<Rule> loadRulesFromDisk() throws IOException {
    List<Rule> rules = new ArrayList<>();
    
    // 获取规则目录路径，可以通过启动参数指定，默认为"rules"
    String rulesDirPath = System.getProperty("rules.dir", "rules");
    Path rulesDir = Paths.get(rulesDirPath);
    if (!Files.exists(rulesDir)) {
      return rules; // Return empty list if rules directory doesn't exist
    }
    
    // 使用Files.walk遍历目录及其子目录
    try (Stream<Path> paths = Files.walk(rulesDir)) {
      paths.filter(Files::isRegularFile) // 只处理文件
        .filter(path -> path.toString().endsWith(".md")) // 只处理.md文件
        .forEach(path -> {
          try {
            List<String> lines = Files.readAllLines(path);
            if (lines.size() >= 2) {
              String name = lines.get(0).trim();
              String description = lines.get(1).trim();
              
              // 创建inputSchema
              Map<String, Object> properties = new HashMap<>();
              Map<String, Object> inputSchema = new HashMap<>();
              inputSchema.put("type", "object");
              inputSchema.put("properties", properties);

              // 缓存工具内容（第3行到最后）
              if (lines.size() > 2) {
                StringBuilder content = new StringBuilder();
                for (int i = 2; i < lines.size(); i++) {
                  content.append(lines.get(i));
                  if (i < lines.size() - 1) {
                    content.append("\n");
                  }
                }
                Rule rule = new Rule(name, description, content.toString());
                rules.add(rule);
              }
            }
          } catch (IOException e) {
            logger.error("Failed to read rule file: {}, error: {}", path, e.getMessage(), e);
          }
        });
    }
    
    return rules;
  }
  
  private static void setupGitPullTask() {
    String allowPull = System.getProperty("allow.pull", "N");
    if(allowPull.equalsIgnoreCase("Y")) {
      // 每隔5分钟执行一次git pull
      vertx.setPeriodic(5 * 60 * 1000, id -> {
        String rulesDirPath = System.getProperty("rules.dir", "rules");
        pullLatestRules(rulesDirPath);
      });
    }
  }

  private static void pullLatestRules(String rulesDir) {
    try {
      logger.info("开始拉取规则目录最新代码: {}", rulesDir);

      // 构建执行git pull的命令
      ProcessBuilder processBuilder = new ProcessBuilder();
      processBuilder.command("git", "pull");
      processBuilder.directory(Paths.get(rulesDir).toFile());

      // 执行命令
      Process process = processBuilder.start();

      // 读取标准输出
      BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
      String line;
      while ((line = reader.readLine()) != null) {
        logger.info("Git pull output: {}", line);
      }

      // 等待命令执行完成
      int exitCode = process.waitFor();
      if (exitCode == 0) {
        logger.info("成功拉取规则目录最新代码");
        // 更新缓存时间，以便下次请求时重新加载规则
        lastCacheUpdate = 0;
      } else {
        logger.error("Git pull执行失败，退出码: {}", exitCode);
        
        // 读取错误输出
        BufferedReader errorReader = new BufferedReader(new InputStreamReader(process.getErrorStream()));
        String errorLine;
        while ((errorLine = errorReader.readLine()) != null) {
          logger.error("Git pull error: {}", errorLine);
        }
      }
    } catch (Exception e) {
      logger.error("执行git pull时发生异常", e);
    }
  }
}
